import cv2
import torch
import sys
import os
import time

# 添加 YOLOv5 路径
sys.path.append('yolov5')


def main():
    print("🎯 纯 YOLOv5 绿色小球检测器")
    print("=" * 30)

    # 检查 YOLOv5 目录
    if not os.path.exists('yolov5'):
        print("❌ YOLOv5 目录不存在")
        print("请先运行:")
        print("git clone https://github.com/ultralytics/yolov5.git")
        print("cd yolov5")
        print("git checkout v7.0")
        return

    try:
        # 使用本地 YOLOv5
        model = torch.hub.load('yolov5', 'custom', path='best.pt', source='local')
        model.conf = 0.5  # 置信度阈值
        print("✅ YOLOv5 模型加载成功")
    except Exception as e:
        print(f"❌ 模型加载失败: {e}")
        return

    # 打开摄像头
    cap = cv2.VideoCapture(3)
    if not cap.isOpened():
        print("❌ 无法打开摄像头")
        return

    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

    print("✅ 摄像头已打开")
    print("按 'q' 退出，按 's' 保存截图")

    frame_count = 0
    start_time = time.time()

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1

        try:
            # YOLOv5 检测
            results = model(frame)

            # 渲染结果
            results.render()
            annotated_frame = results.imgs[0]

            # 获取检测结果
            detections = results.pandas().xyxy[0]
            ball_count = len(detections)

            # 计算FPS
            elapsed_time = time.time() - start_time
            fps = frame_count / elapsed_time if elapsed_time > 0 else 0

            # 显示信息
            cv2.putText(annotated_frame, f"FPS: {fps:.1f}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.putText(annotated_frame, f"Balls: {ball_count}", (10, 60),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

            cv2.imshow('YOLOv5 Detection', annotated_frame)

        except Exception as e:
            print(f"检测错误: {e}")
            cv2.imshow('YOLOv5 Detection', frame)

        # 处理按键
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            filename = f"capture_{timestamp}.jpg"
            cv2.imwrite(filename, annotated_frame)
            print(f"📷 保存: {filename}")

    cap.release()
    cv2.destroyAllWindows()
    print("✅ 检测结束")


if __name__ == "__main__":
    main()